
import os
from binascii import b2a_hex

from pdfminer.pdfparser import PDFParser
from pdfminer.pdfdocument import PDFDocument
from pdfminer.pdfinterp import PDFPageInterpreter
from pdfminer.pdfinterp import PDFResourceManager
from pdfminer.converter import PDFPageAggregator
from pdfminer.layout import LAParams, LTImage
from pdfminer.pdfpage import PDFPage

'''pdfminer无法获取PDF中的图片'''

class PDFBase(object):
    def __init__(self, filepath):
        fp = open(filepath, 'rb')
        self.pdfparser = PDFParser(fp)
        self.doc = PDFDocument(self.pdfparser)
        self.resource = PDFResourceManager()
        self.laparm = LAParams()
        self.device = PDFPageAggregator(self.resource, laparams=self.laparm)
        self.interpreter = PDFPageInterpreter(self.resource, self.device)

    def get_text(self, start_page, end_page):
        '''获得文字'''
        outputs = []
        for i, page in enumerate(PDFPage.create_pages(self.doc)):
            if start_page <= i + 1 and i + 1 <= end_page:
                self.interpreter.process_page(page)
                layout = self.device.get_result()
                for out in layout:
                    if hasattr(out, 'get_text'): 
                        outputs.append(out.get_text())
        return outputs

    def write_file(self, folder, filename, filedata, flags='w'):
        """
        :param flags:写入文件格式(w,wb,a) floder
        将文件数据写入文件夹和文件名组合
        :return
        """
        try:
            if not os.path.exists(folder):
                os.makedirs(folder)
            with open(os.path.join(folder, filename), flags) as f:
                f.write(filedata)
        except:
            return False
        return True

    def determine_image_type(self,stream_first_4_bytes):
        """Find out the image file type based on the magic number comparison of the first 4 (or 2) bytes"""
        '''
        1. 首先以二进制方式编辑这个文件
        vi -b datafile

        2. 使用 xxd 转换为 16 进制
        :%!xxd
        '''
        file_type = None
        # b2a_hex 返回二进制数据对应的16进制表示
        bytes_as_hex = b2a_hex(stream_first_4_bytes)
        if bytes_as_hex.startswith('ffd8'):
            file_type = '.jpeg'
        elif bytes_as_hex == '89504e47':
            file_type = '.png'
        elif bytes_as_hex == '47494638':
            file_type = '.gif'
        elif bytes_as_hex.startswith('424d'):
            file_type = '.bmp'
        return file_type

    def save_image(self, out, dir):
        try:
            result = None
            if out.stream:
                file_stream = out.stream.get_rawdata()
                file_ext = self.determine_image_type(file_stream[0:4])
            if file_ext:
                file_name = ''.join([str(page_number), '_', out.name, file_ext])
            if self.write_file(images_folder, file_name, out.stream.get_rawdata(), flags='wb'):
                result = file_name
            print('success')
            return result
        except:
            pass
            # print('error')

    def get_image(self, start_page, end_page, imgpath):
        '''获得图片'''
        for i, page in enumerate(PDFPage.create_pages(self.doc)):
            if start_page <= i + 1 and i + 1 <= end_page:
                self.interpreter.process_page(page)
                layout = self.device.get_result()
                for out in layout:
                    self.save_image(out, imgpath)
                        

class MyPDF(PDFBase):
    def do_something(self):
        # 获取第1页到第2页的图片
        self.get_image(1, 2, './img')

pdf = MyPDF('图片.pdf')
pdf.do_something()
